import csv
import re
import pandas as pd
import sqlite3
import logging
from utils.instances import TOKENIZER, LLM
from utils import prompts
from langchain_core.prompts import ChatPromptTemplate
import utils.configFinRAG as configFinRAG

# 配置日志记录
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# 全局变量，用于存储示例数据
g_example_question_list = []
g_example_sql_list = []
g_example_info_list = []
g_example_fa_list = []
g_example_token_list = []


# 封装 Jaccard 相似度计算函数
def jaccard_similarity(list1, list2):
    try:
        intersection = len(set(list1) & set(list2))
        union = len(set(list1)) + len(set(list2))
        return intersection / union if union != 0 else 0
    except Exception as e:
        logging.error(f"Jaccard 相似度计算出错: {e}")
        return 0


def generate_sql(user_question, language_model, example_questions, example_sqls, example_token_lists, num_examples=5):
    try:
        # 匹配基金代码，一般基金代码是6位数字
        fund_code_pattern = r'\d{6}'
        sql_start_pattern = '```sql'
        sql_end_pattern = '```'
        temp_question = user_question
        # 找出问题里的基金代码
        fund_code_list = re.findall(fund_code_pattern, temp_question)
        processed_question = temp_question
        # 移除问题中的基金代码，避免影响相似度计算
        for fund_code in fund_code_list:
            processed_question = processed_question.replace(fund_code, ' ')
        question_tokens = TOKENIZER(processed_question)['input_ids']
        # 计算问题与示例问题的 Jaccard 相似度
        similarities = [jaccard_similarity(question_tokens, token_list) for token_list in example_token_lists]
        # 按相似度从高到低排序，取前 num_examples 个
        top_indices = sorted(range(len(similarities)), key=lambda i: similarities[i], reverse=True)[:num_examples]
        prompt_length = 0
        selected_indices = []
        # 根据提示长度筛选合适的示例问题
        for index in top_indices:
            prompt_length += len(example_questions[index]) + len(example_sqls[index])
            if prompt_length > 2000:
                break
            selected_indices.append(index)
        # 创建聊天提示模板
        prompt_template = ChatPromptTemplate.from_template(prompts.GENERATE_SQL_TEMPLATE)
        # 拼接选中的示例问题和对应的 SQL 语句
        example_prompts = '\n'.join([
            f"问题：{example_questions[index]}\nSQL：{example_sqls[index]}"
            for index in selected_indices
        ])
        # 构建提示链并调用语言模型生成 SQL
        chain = prompt_template | language_model
        response = chain.invoke(
            {"examples": example_prompts, "table_info": prompts.TABLE_INFO, "question": temp_question})
        sql = response.content
        # 提取生成的 SQL 语句
        start_index = sql.find(sql_start_pattern) + len(sql_start_pattern)
        end_index = sql[start_index:].find(sql_end_pattern) + start_index if start_index >= 0 else -1
        if start_index < end_index:
            sql = sql[start_index:end_index]
            return prompt_template.invoke({"examples": example_prompts, "table_info": prompts.TABLE_INFO,
                                           "question": temp_question}), sql
        else:
            logging.error(f"生成 SQL 出错: {user_question}")
            return "error", "error"
    except Exception as e:
        logging.error(f"generate_sql 函数出错: {e}")
        return "error", "error"


def query_db(sql, cursor):
    try:
        cursor.execute(sql)
        result = cursor.fetchall()
        return True, result
    except Exception as e:
        logging.error(f"查询数据库出错: {sql}, 错误信息: {e}")
        return False, []


def generate_answer(query, exc_result, language_model, example_questions, example_info_list, example_fa_list,
                    example_token_lists):
    # 这里简单示例返回查询结果，可根据实际需求完善生成答案的逻辑
    if exc_result:
        return str(exc_result)
    else:
        return "未找到相关结果"


def sql_retrieve_chain(query):
    global g_example_question_list, g_example_sql_list, g_example_info_list, g_example_fa_list, g_example_token_list
    if len(g_example_question_list) <= 0:
        try:
            sql_examples_file = pd.read_csv(configFinRAG.sql_examples_path, delimiter=",", header=0)
            for index in range(len(sql_examples_file)):
                g_example_question_list.append(sql_examples_file.iloc[index]['问题'])
                g_example_sql_list.append(sql_examples_file.iloc[index]['SQL'])
                g_example_info_list.append(sql_examples_file.iloc[index]['资料'])
                g_example_fa_list.append(sql_examples_file.iloc[index]['FA'])
                tokens = TOKENIZER(sql_examples_file.iloc[index]['问题'])
                tokens = tokens['input_ids']
                g_example_token_list.append(tokens)
        except Exception as e:
            logging.error(f"读取示例数据出错: {e}")
            return "读取示例数据失败，请检查配置和文件"

    result_prompt, sql = generate_sql(query, LLM, g_example_question_list, g_example_sql_list, g_example_token_list)
    if sql == "error":
        return "生成 SQL 语句失败，请检查问题输入"

    try:
        conn = sqlite3.connect(configFinRAG.db_path)  # 假设在 configFinRAG 中配置了数据库路径
        cs = conn.cursor()
        success_flag, exc_result = query_db(sql, cs)
        conn.close()
    except Exception as e:
        logging.error(f"数据库连接或查询出错: {e}")
        return "数据库连接或查询出错，请检查数据库配置"

    answer = generate_answer(query, exc_result, LLM, g_example_question_list, g_example_info_list, g_example_fa_list,
                             g_example_token_list)
    return answer


if __name__ == '__main__':
    user_query = "查询基金代码 006571 的基金经理是谁"
    result = sql_retrieve_chain(user_query)
    print(result)
